import fasttext as ft
import pickle


MODELS = {
    'ft_title': '/data/data/new_class1_ft_title.bin',
    'ft_content': '/data/data/new_class1_ft_content.bin',
    'lg_title': '/data/data/new_class1_ft_title.tsv_CV.pkl,/data/data/new_class1_ft_title.tsv_TFIDF.pkl,/data/data/new_class1_ft_title.tsv_LG.pkl',
    'lg_content': '/data/data/new_class1_ft_content.tsv_CV.pkl,/data/data/new_class1_ft_content.tsv_TFIDF.pkl,/data/data/new_class1_ft_content.tsv_LG.pkl',
    'nb_title': '/data/data/new_class1_ft_title.tsv_CV.pkl,/data/data/new_class1_ft_title.tsv_TFIDF.pkl,/data/data/new_class1_ft_title.tsv_NB.pkl',
    'nb_content': '/data/data/new_class1_ft_content.tsv_CV.pkl,/data/data/new_class1_ft_content.tsv_TFIDF.pkl,/data/data/new_class1_ft_content.tsv_LG.pkl'
}


class PickledClassifierModel:
    def __init__(self, pfiles):
        model_files = [k for k in pfiles.split(',') if k != '']
        self.models = [pickle.load(open(k, 'rb')) for k in model_files]

    def predict(self, text: list):
        x = text
        for m in self.models[0: -1]:
            x = m.transform(x)

        x = self.models[-1].predict(x)
        return [x]


def load_models():
    ft_title_model = ft.load_model(MODELS['ft_title'])
    ft_content_model = ft.load_model(MODELS['ft_content'])
    lg_title_model = PickledClassifierModel(MODELS['lg_title'])
    lg_content_model = PickledClassifierModel(MODELS['lg_content'])
    nb_title_model = PickledClassifierModel(MODELS['nb_title'])
    nb_content_model = PickledClassifierModel(MODELS['nb_content'])

    return [ft_title_model, lg_title_model, nb_title_model], [ft_content_model, lg_content_model,  nb_content_model]


def predict(models, texts):
    return [k.predict(texts) for k in models]


if __name__ == '__main__':
    title_models, content_models = load_models()
    target_title_content_file = '/data/data/new_class1_ft_title_content.tsv'
    out_file = '/data/data/new_class1_fusion.tsv'
    out_fd = open(out_file, 'w+')
    # out = predict(models[0], ['吉利 旗舰 SUV ， 配 沃尔沃 动力 ， 值得 买 吗 ？'])
    for l in open(target_title_content_file):
        d = l.strip().split('\t')
        if len(d) != 3:
            continue

        target, title, content = d

        target = target.replace('__label__', '')
        if title.strip() != '':
            title_predicts = predict(title_models, [title])
        else:
            title_predicts = [[['<UNK_CLASS>']] for k in title_models]
        if content.strip() != '':
            content_predicts = predict(content_models, [content])
        else:
            content_predicts = [[['<UNK_CLASS>']] for k in content_models]

        title_predicts = [k[0][0].replace('__label__', '') for k in title_predicts]
        content_predicts = [k[0][0].replace('__label__', '') for k in content_predicts]
        out = target + '\t' + ' '.join(title_predicts + content_predicts)
        out_fd.write(out + '\n')


